# -*- coding: utf-8 -*-
from .db import get_db
import datetime
import pandas as pd
import numpy as np

# 获取stocks数据的api
def get_stock_data(code, **params):
    if 'start' in params:
        start = params['start']
        if 'end' in params:
            end = params['end']
        else:
            end = datetime.datetime.now().strftime('%Y-%m-%d')
    else:
        if 'end' in params:
            end = params['end']
            end_date = pd.Timestamp(end)
            start_date = end_date - datetime.timedelta(days=90)
            start = start_date.strftime('%Y-%m-%d')
        else:
            end_date = datetime.datetime.now()
            start_date = end_date - datetime.timedelta(days=90)
            start = start_date.strftime('%Y-%m-%d')
            end = end_date.strftime('%Y-%m-%d')

    data = get_db().stocks.find({'code': code, 'date': {'$gte': start, '$lte': end}})
    data = pd.DataFrame(list(data))
    data = data.loc[:, ['date', 'code', 'open', 'low', 'high', 'close', 'volume']]

    return data

# 搜索股票数据
def search_stocks_data(param):
    return get_db().stocks.find(param)

# 获取最后的交易日期
def get_last_trading_day():
    data = get_db().stocks.find({"code": "hs300"}).sort([("date", -1)]).limit(1)
    data = list(data)
    lastday = data[0]['date']
    # print("lastday = ", lastday)
    return lastday

# 获取股票的基础数据的api, 返回dateframe
def get_stocks_basics(date=None):
    if date == None:
        date = get_last_trading_day()

    data = get_db().stocks_basics.find({"date": date})
    data = pd.DataFrame(list(data))
    del data['_id']
    return data

# 搜索股票的基础数据的api
def search_stocks_basics(param):
    return get_db().stocks_basics.find(param)